#script to perform intensity interpolation and sum 
#based on distance of illumination foci to each pixel

import os, sys
import numpy, pylab
from mpl_toolkits.mplot3d import Axes3D
from scipy.interpolate import griddata

def destripify(filename, shape=(250, 1024, 1024), display=True):
    basename, ext = os.path.splitext(filename)
    print "Loading data file:", filename
    data1 = numpy.fromfile(filename, dtype=numpy.float64).reshape(shape)
    print "Loading data file:", basename + '_x' + ext
    dataX = numpy.fromfile(basename + '_x' + ext, dtype=numpy.float64).reshape(shape)
    print "Loading data file:", basename + '_y' + ext
    dataY = numpy.fromfile(basename + '_y' + ext, dtype=numpy.float64).reshape(shape)
    print "Done loading"

    final_image = numpy.zeros(data1.shape[1:], dtype=numpy.float64)

    for x_d in range(data1.shape[1]):
        print "Calculating integral for row", x_d
        for y_d in range(data1.shape[2]):
            data_intensity_z = data1[:, x_d, y_d]
            x_shift = dataX[:, x_d, y_d]
            y_shift = dataY[:, x_d, y_d]
            nonzero = data_intensity_z != 0
            data_intensity_z = data_intensity_z[nonzero]
            x_shift = x_shift[nonzero]
            y_shift = y_shift[nonzero]

            if len(data_intensity_z) > 3:
                xi = numpy.linspace(-5, 5, 30)
                yi = numpy.linspace(-5, 5, 30)
                try:
                    zi = griddata((x_shift, y_shift), data_intensity_z,
                                  (xi[None, :], yi[:, None]),
                                  method='cubic', fill_value=0)
                except RuntimeError:
                    print "QHull error"
                    continue
                final_image[x_d, y_d] = zi.mean()

                if display:
                    x_coords, y_coords = numpy.meshgrid(xi, yi)
                    fig = pylab.figure()
                    fig2 = pylab.figure()
                    ax = fig.add_subplot(111, projection='3d')
                    ax2 = fig2.add_subplot(111)
                    for c, m in [('r', 'o')]:
                        ax.scatter(x_shift, y_shift, data_intensity_z, c=c, marker=m)
                        ax2.scatter(x_shift, y_shift, c=c, marker=m)
                    for c, m in [('b', '.')]:
                        ax.scatter(x_coords, y_coords, zi, c=c, marker=m)
                        ax2.scatter(x_coords, y_coords, c=c, marker=m)
                    ax.set_xlabel('X Label')
                    ax.set_ylabel('Y Label')
                    ax.set_zlabel('Z Label')
                    fig.show()
                    fig2.show()
                    raw_input()
                    pylab.close(fig)
                    pylab.close(fig2)
    final_image.tofile(basename + '_destriped' + ext)

def process_folder(shape=(250, 1024, 1024)):
    import Tkinter, tkFileDialog, tkSimpleDialog, glob

    tkroot = Tkinter.Tk()
    tkroot.withdraw()
    data_filename = str(os.path.normpath(tkFileDialog.askopenfilename(
        title='Select one of your enderlein "_frames" images',
        filetypes=[('Raw binary', '.raw')],
        defaultextension='.raw',
        initialdir=os.getcwd()
        )))
    data_dir = os.path.dirname(data_filename)

    while True:
        wildcard_data_filename = tkSimpleDialog.askstring(
            title='Filename pattern',
            prompt=("Use '?' as a wildcard"),
            initialvalue=os.path.split(data_filename)[1])
        data_filenames_list = sorted(glob.glob(
            os.path.join(data_dir, wildcard_data_filename)))
        print "Data filenames:"
        for f in data_filenames_list:
            print '.   ' + f
        response = raw_input('Are these the files you want to process? [y]/n:')
        if response == 'n':
            continue
        else:
            break

    for f in data_filenames_list:
        destripify(f, shape=shape, display=False)

process_folder(shape=(250, 1024, 1024))

